import torch
from torch import nn
import torch.nn.functional as F

class MultiFocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, inputs, targets):
        BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)
        
        # 自动计算alpha
        if self.alpha == 'auto':
            pos_rate = targets.mean(dim=0)
            alpha = 1 / (pos_rate + 1e-8)
            alpha = alpha / alpha.mean()  # 归一化
        else:
            alpha = self.alpha
            
        alpha_factor = alpha * targets + (1 - alpha) * (1 - targets)
        focal_loss = alpha_factor * (1 - pt)**self.gamma * BCE_loss
        
        return focal_loss.mean()
